import tensorflow as tf
import os

VOC_dir = 'data/VOCdevkit/VOC2012'

AUTOTUNE = tf.data.experimental.AUTOTUNE


def preprocess_image(image):
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, [128, 128])
    image /= 255.0  # normalize to [0,1] range

    return image


def load_and_preprocess_image(path):
    image = tf.io.read_file(path)
    return preprocess_image(image)


def load_and_preprocess_image_label(feature, label):
    return load_and_preprocess_image(feature), load_and_preprocess_image(label)


# 读取训练数据集
def read_Voc(file_type='train'):
    txt_name = os.path.join(VOC_dir, 'ImageSets', 'Segmentation', file_type + '.txt')

    with open(txt_name, 'r') as f:
        images = f.read().split()

    features_path, labels_path = [], []
    for i, fname in enumerate(images):
        feature_path = os.path.join(VOC_dir, 'JPEGImages', f'{fname}.jpg')
        features_path.append(feature_path)
        label_path = os.path.join(VOC_dir, 'SegmentationClass', f'{fname}.png')
        labels_path.append(label_path)

    dataset = tf.data.Dataset.from_tensor_slices((features_path, labels_path))
    dataset = dataset.map(load_and_preprocess_image_label, num_parallel_calls=AUTOTUNE)
    return dataset

# train_dataset = read_Voc()
# val_dataset = read_Voc(file_type='val')
